from flask import Flask, Response, json,request
from dashscope import Generation
from dashscope.api_entities.dashscope_response import Role
from flask_cors import CORS
import dashscope

# import numpy as np

app = Flask(__name__)
CORS(app, resources="/*")

# 接收全部的信息
messages = []

@app.route('/llm/request')
def stream_numbers():
    dashscope.api_key = 'sk-29ba8d911efb41dfa6092d07fe9c8e5f'
    global messages
    query = request.args.get('query', default='default query')

    def chat():
        print(query)
        messages.append({'role': Role.USER, 'content': query})
        whole_message = ''
        responses = Generation.call(Generation.Models.qwen_turbo, messages=messages, result_format='message', stream=True,
                                    incremental_output=True)

        for response in responses:
            part_message = response.output.choices[0]['message']['content']
            whole_message += part_message
            print(part_message, end='')
            json_data = json.dumps({"message": response.output.choices[0]['message']['content']})
            yield f"data: {json_data}\n\n"  # 按照SSE格式发送数据

        messages.append({'role': 'assistant', 'content': whole_message})
        json_data = json.dumps({"message": 'done'})
        yield f"data: {json_data}\n\n"  # 按照SSE格式发送数据
        print('结束')
    headers = {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'X-Accel-Buffering': 'no',
    }

    return Response(chat(), content_type='text/event-stream', headers=headers)


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
